#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <pthread.h>
#include <netdb.h>
#include <signal.h>

#define MAX_BUFFER 1024
#define MAX_NAME 50
#define MAX_GROUP 20
#define MAX_CLIENTS 100
#define MAX_GROUPS 50

// 消息类型定义 (与客户端一致)
#define REGISTER 1
#define LOGIN 2
#define LOGOUT 3
#define LIST_USERS 4
#define LIST_GROUPS 5
#define PRIVATE_MSG 6
#define CREATE_GROUP 7
#define JOIN_GROUP 8
#define GROUP_MSG 9
#define LEAVE_GROUP 10
#define BROADCAST 11

typedef struct {
    int type;
    char sender[MAX_NAME];
    char target[MAX_NAME];
    char message[MAX_BUFFER];
} ChatMessage;

typedef struct {
    char username[MAX_NAME];
    char password[MAX_NAME];
    int online;
    int sockfd;
    pthread_t thread_id;
} User;

typedef struct {
    char groupname[MAX_NAME];
    char members[MAX_CLIENTS][MAX_NAME];
    int member_count;
} Group;

User users[MAX_CLIENTS];
Group groups[MAX_GROUPS];
int user_count = 0;
int group_count = 0;
pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;

void broadcast_message(const char *message, int exclude_sock);
void send_to_group(const char *groupname, const char *message, const  char *sender);
void *client_handler(void *arg);
void save_user_data();
void load_user_data();

// 查找用户
User *find_user(const char *username) {
    for (int i = 0; i < user_count; i++) {
        if (strcmp(users[i].username, username) == 0) {
            return &users[i];
        }
    }
    return NULL;
}

// 查找群组
Group *find_group(const char *groupname) {
    for (int i = 0; i < group_count; i++) {
        if (strcmp(groups[i].groupname, groupname) == 0) {
            return &groups[i];
        }
    }
    return NULL;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "用法: %s <监听IP> <端口号>\n", argv[0]);
        exit(1);
    }

    char *host = argv[1];
    int port = atoi(argv[2]);

    // 加载用户数据
    load_user_data();

    // 创建UDP套接字用于注册
    int udp_sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_sock < 0) {
        perror("UDP套接字创建失败");
        exit(1);
    }

    // 创建TCP套接字用于客户端连接
    int tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (tcp_sock < 0) {
        perror("TCP套接字创建失败");
        close(udp_sock);
        exit(1);
    }

    // 设置地址重用
    int opt = 1;
    setsockopt(tcp_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    setsockopt(udp_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    // 设置服务器地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    server_addr.sin_addr.s_addr = inet_addr(host);

    // 绑定套接字
    if (bind(udp_sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("UDP绑定失败");
        close(udp_sock);
        close(tcp_sock);
        exit(1);
    }

    if (bind(tcp_sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("TCP绑定失败");
        close(udp_sock);
        close(tcp_sock);
        exit(1);
    }

    // 监听TCP连接
    if (listen(tcp_sock, 10) < 0) {
        perror("监听失败");
        close(udp_sock);
        close(tcp_sock);
        exit(1);
    }

    printf("服务器已启动，监听 %s:%d\n", host, port);

    // 设置文件描述符集合
    fd_set read_fds;
    int max_fd = (udp_sock > tcp_sock) ? udp_sock : tcp_sock;
    max_fd = (max_fd > STDIN_FILENO) ? max_fd : STDIN_FILENO;

    while (1) {
        FD_ZERO(&read_fds);
        FD_SET(udp_sock, &read_fds);
        FD_SET(tcp_sock, &read_fds);
        FD_SET(STDIN_FILENO, &read_fds);

        if (select(max_fd + 1, &read_fds, NULL, NULL, NULL) < 0) {
            perror("select错误");
            continue;
        }

        // 处理UDP注册请求
        if (FD_ISSET(udp_sock, &read_fds)) {
            struct sockaddr_in client_addr;
            socklen_t client_len = sizeof(client_addr);
            
            ChatMessage reg_msg;
            recvfrom(udp_sock, &reg_msg, sizeof(reg_msg), 0, 
                     (struct sockaddr*)&client_addr, &client_len);
            
            if (reg_msg.type == REGISTER) {
                ChatMessage resp;
                resp.type = REGISTER;
                strncpy(resp.sender, reg_msg.sender, MAX_NAME);
                
                pthread_mutex_lock(&lock);
                User *user = find_user(reg_msg.sender);
                
                if (user) {
                    strcpy(resp.message, "用户名已存在");
                } else if (user_count >= MAX_CLIENTS) {
                    strcpy(resp.message, "用户数量已达上限");
                } else {
                    strcpy(users[user_count].username, reg_msg.sender);
                    strcpy(users[user_count].password, reg_msg.message);
                    users[user_count].online = 0;
                    users[user_count].sockfd = -1;
                    user_count++;
                    
                    strcpy(resp.message, "SUCCESS");
                    save_user_data();
                    printf("新用户注册: %s\n", reg_msg.sender);
                }
                pthread_mutex_unlock(&lock);
                
                sendto(udp_sock, &resp, sizeof(resp), 0, 
                       (struct sockaddr*)&client_addr, client_len);
            }
        }

        // 处理新的TCP连接
        if (FD_ISSET(tcp_sock, &read_fds)) {
            struct sockaddr_in client_addr;
            socklen_t client_len = sizeof(client_addr);
            
            int client_sock = accept(tcp_sock, (struct sockaddr*)&client_addr, &client_len);
            if (client_sock < 0) {
                perror("接受连接失败");
                continue;
            }
            
            // 创建线程处理客户端
            pthread_t thread;
            int *sock_ptr = (int*)malloc(sizeof(int)); // 修复内存分配
            *sock_ptr = client_sock;
            
            if (pthread_create(&thread, NULL, client_handler, (void*)sock_ptr)) {
                perror("线程创建失败");
                close(client_sock);
                free(sock_ptr);
            }
        }

        // 处理管理员输入
        if (FD_ISSET(STDIN_FILENO, &read_fds)) {
            char input[MAX_BUFFER];
            fgets(input, MAX_BUFFER, stdin);
            input[strcspn(input, "\n")] = '\0';
            
            if (strlen(input) == 0) continue;
            
            char *cmd = strtok(input, " ");
            if (cmd == NULL) continue;
            
            if (strcmp(cmd, "broadcast") == 0) {
                char *message = strtok(NULL, "");
                if (message) {
                    ChatMessage msg;
                    msg.type = BROADCAST;
                    strcpy(msg.sender, "管理员");
                    strcpy(msg.message, message);
                    
                    pthread_mutex_lock(&lock);
                    for (int i = 0; i < user_count; i++) {
                        if (users[i].online) {
                            send(users[i].sockfd, &msg, sizeof(msg), 0);
                        }
                    }
                    pthread_mutex_unlock(&lock);
                }
            }
            else if (strcmp(cmd, "sendto") == 0) {
                char *target = strtok(NULL, " ");
                char *message = strtok(NULL, "");
                
                if (target && message) {
                    pthread_mutex_lock(&lock);
                    User *user = find_user(target);
                    if (user && user->online) {
                        ChatMessage msg;
                        msg.type = PRIVATE_MSG;
                        strcpy(msg.sender, "管理员");
                        strcpy(msg.target, target);
                        strcpy(msg.message, message);
                        send(user->sockfd, &msg, sizeof(msg), 0);
                    }
                    pthread_mutex_unlock(&lock);
                }
            }
            else if (strcmp(cmd, "sendgroup") == 0) {
                char *group = strtok(NULL, " ");
                char *message = strtok(NULL, "");
                
                if (group && message) {
                    pthread_mutex_lock(&lock);
                    Group *grp = find_group(group);
                    if (grp) {
                        ChatMessage msg;
                        msg.type = GROUP_MSG;
                        strcpy(msg.sender, "管理员");
                        strcpy(msg.target, group);
                        strcpy(msg.message, message);
                        
                        for (int i = 0; i < grp->member_count; i++) {
                            User *user = find_user(grp->members[i]);
                            if (user && user->online) {
                                send(user->sockfd, &msg, sizeof(msg), 0);
                            }
                        }
                    }
                    pthread_mutex_unlock(&lock);
                }
            }
            else if (strcmp(cmd, "users") == 0) {
                printf("\n=== 用户列表 ===\n");
                pthread_mutex_lock(&lock);
                for (int i = 0; i < user_count; i++) {
                    printf("%s - %s\n", users[i].username, 
                           users[i].online ? "在线" : "离线");
                }
                pthread_mutex_unlock(&lock);
            }
            else if (strcmp(cmd, "groups") == 0) {
                printf("\n=== 群组列表 ===\n");
                pthread_mutex_lock(&lock);
                for (int i = 0; i < group_count; i++) {
                    printf("%s (%d 名成员)\n", groups[i].groupname, groups[i].member_count);
                }
                pthread_mutex_unlock(&lock);
            }
            else {
                printf("可用命令:\n");
                printf("  broadcast <消息>       - 广播消息\n");
                printf("  sendto <用户> <消息>   - 发送私信\n");
                printf("  sendgroup < 组> <消息> - 发送群消息\n");
                printf("  users                  - 列出所有用户\n");
                printf("  groups                 - 列出所有群组\n");
            }
        }
    }

    close(udp_sock);
    close(tcp_sock);
    return 0;
}

void *client_handler(void *arg) {
    int sockfd = *((int*)arg);
    free(arg);
    
    ChatMessage login_msg;
    if (recv(sockfd, &login_msg, sizeof(login_msg), 0) <= 0) {
        close(sockfd);
        return NULL;
    }
    
    if (login_msg.type != LOGIN) {
        close(sockfd);
        return NULL;
    }
    
    char *username = login_msg.sender;
    char *password = login_msg.message;
    
    ChatMessage resp;
    resp.type = LOGIN;
    strcpy(resp.sender, username);
    
    pthread_mutex_lock(&lock);
    User *user = find_user(username);
    
    if (!user) {
        strcpy(resp.message, "用户不存在");
        send(sockfd, &resp, sizeof(resp), 0);
        pthread_mutex_unlock(&lock);
        close(sockfd);
        return NULL;
    }
    
    if (strcmp(user->password, password) != 0) {
        strcpy(resp.message, "密码错误");
        send(sockfd, &resp, sizeof(resp), 0);
        pthread_mutex_unlock(&lock);
        close(sockfd);
        return NULL;
    }
    
    if (user->online) {
        strcpy(resp.message, "用户已在线");
        send(sockfd, &resp, sizeof(resp), 0);
        pthread_mutex_unlock(&lock);
        close(sockfd);
        return NULL;
    }
    
    // 登录成功
    user->online = 1;
    user->sockfd = sockfd;
    strcpy(resp.message, "SUCCESS");
    send(sockfd, &resp, sizeof(resp), 0);
    
    // 广播登录通知
    ChatMessage notify;
    notify.type = BROADCAST;
    sprintf(notify.message, "%s 已上线", username);
    
    // 解锁后再广播，避免死锁
    pthread_mutex_unlock(&lock);
    broadcast_message(notify.message, sockfd);
    
    printf("%s 已登录\n", username);
    
    // 处理客户端请求
    while (1) {
        ChatMessage msg;
        ssize_t n = recv(sockfd, &msg, sizeof(msg), 0);
        if (n <= 0) break;
        
        pthread_mutex_lock(&lock);
        User *sender = find_user(username);
        if (!sender || !sender->online) {
            pthread_mutex_unlock(&lock);
            break;
        }
        
        // 处理消息
        memset(&resp, 0, sizeof(resp));
        resp.type = msg.type;
        strcpy(resp.sender, username);
        
        switch (msg.type) {
            case LIST_USERS: {
                resp.message[0] = '\0';
                for (int i = 0; i < user_count; i++) {
                    if (users[i].online) {
                        strcat(resp.message, users[i].username);
                        strcat(resp.message, "\n");
                    }
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case LIST_GROUPS: {
                resp.message[0] = '\0';
                for (int i = 0; i < group_count; i++) {
                    strcat(resp.message, groups[i].groupname);
                    strcat(resp.message, "\n");
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case PRIVATE_MSG: {
                User *receiver = find_user(msg.target);
                if (receiver && receiver->online) {
                    ChatMessage private_msg;
                    private_msg.type = PRIVATE_MSG;
                    strcpy(private_msg.sender, username);
                    strcpy(private_msg.target, msg.target);
                    strcpy(private_msg.message, msg.message);
                    
                    send(receiver->sockfd, &private_msg, sizeof(private_msg), 0);
                    strcpy(resp.message, "消息已发送");
                } else {
                    strcpy(resp.message, "用户不在线或不存在");
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case CREATE_GROUP: {
                Group *existing = find_group(msg.target);
                if (existing) {
                    strcpy(resp.message, "群组已存在");
                } else if (group_count >= MAX_GROUPS) {
                    strcpy(resp.message, "群组数量已达上限");
                } else {
                    Group new_group;
                    strcpy(new_group.groupname, msg.target);
                    new_group.member_count = 0;
                    
                    groups[group_count] = new_group;
                    group_count++;
                    
                    strcpy(resp.message, "群组创建成功");
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case JOIN_GROUP: {
                Group *group = find_group(msg.target);
                if (!group) {
                    strcpy(resp.message, "群组不存在");
                } else {
                    // 检查是否已在群组中
                    int in_group = 0;
                    for (int i = 0; i  < group->member_count; i++) {
                        if (strcmp(group->members[i], username) == 0) {
                            in_group = 1;
                            break;
                        }
                    }
                    
                    if (in_group) {
                        strcpy(resp.message, "你已在此群组中");
                    } else {
                        strcpy(group->members[group->member_count], username);
                        group->member_count++;
                        strcpy(resp.message, "已加入群组");
                    }
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case GROUP_MSG: {
                Group *group = find_group(msg.target);
                if (!group) {
                    strcpy(resp.message, "群组不存在");
                    send(sockfd, &resp, sizeof(resp), 0);
                } else {
                    // 检查用户是否在群组中
                    int in_group = 0;
                    for (int i = 0; i < group->member_count; i++) {
                        if (strcmp(group->members[i], username) == 0) {
                            in_group = 1;
                            break;
                        }
                    }
                    
                    if (!in_group) {
                        strcpy(resp.message, "你不在该群组中");
                        send(sockfd, &resp, sizeof(resp), 0);
                    } else {
                        // 发送消息给所有群组成员（除了发送者自己）
                        for (int i = 0; i < group->member_count; i++) {
                            if (strcmp(group->members[i], username) == 0) continue;
                            
                            User *member = find_user(group->members[i]);
                            if (member && member->online) {
                                ChatMessage group_msg;
                                group_msg.type = GROUP_MSG;
                                strcpy(group_msg.sender, username);
                                strcpy(group_msg.target, group->groupname);
                                strcpy(group_msg.message, msg.message);
                                send(member->sockfd, &group_msg, sizeof(group_msg), 0);
                            }
                        }
                        strcpy(resp.message, "消息已发送");
                        send(sockfd, &resp, sizeof(resp), 0);
                    }
                }
                break;
            }
            
            case LEAVE_GROUP: {
                Group *group = find_group(msg.target);
                if (!group) {
                    strcpy(resp.message, "群组不存在");
                } else {
                    // 从群组中移除用户
                    int found = 0;
                    for (int i = 0; i < group->member_count; i++) {
                        if (strcmp(group->members[i], username) == 0) {
                            found = 1;
                            // 移动数组元素
                            for (int j = i; j < group->member_count - 1; j++) {
                                strcpy(group->members[j], group->members[j + 1]);
                            }
                            group->member_count--;
                            break;
                        }
                    }
                    
                    if (found) {
                        strcpy(resp.message, "已离开群组");
                    } else {
                        strcpy(resp.message, "你不在该群组中");
                    }
                }
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case BROADCAST: {
                ChatMessage broadcast_msg;
                broadcast_msg.type = BROADCAST;
                strcpy(broadcast_msg.sender, username);
                strcpy(broadcast_msg.message, msg.message);
                
                for (int i = 0; i < user_count; i++) {
                    if (users[i].online && users[i].sockfd != sockfd) {
                        send(users[i].sockfd, &broadcast_msg, sizeof(broadcast_msg), 0);
                    }
                }
                strcpy(resp.message, "广播已发送");
                send(sockfd, &resp, sizeof(resp), 0);
                break;
            }
            
            case LOGOUT: {
                // 设置用户状态
                user->online = 0;
                user->sockfd = -1;
                pthread_mutex_unlock(&lock); // 解锁后再广播
                
                // 广播登出通知
                char notify_msg[MAX_BUFFER];
                sprintf(notify_msg, "%s 已下线", username);
                broadcast_message(notify_msg, sockfd);
                
                printf("%s 已登出\n", username);
                close(sockfd);
                return NULL;
            }
            
            default:
                strcpy(resp.message, "未知命令");
                send(sockfd, &resp, sizeof(resp), 0);
        }
        
        pthread_mutex_unlock(&lock);
    }
    
    // 客户端断开连接
    pthread_mutex_lock(&lock);
    User *current_user = find_user(username);
    if (current_user) {
        current_user->online = 0;
        current_user->sockfd = -1;
    }
    pthread_mutex_unlock(&lock);
    
    if (current_user) {
        // 广播断开连接通知
        char notify_msg[MAX_BUFFER];
        sprintf(notify_msg, "%s 已断开连接", username);
        broadcast_message(notify_msg, sockfd);
    }
    
    close(sockfd);
    printf("%s 断开连接\n", username);
    return NULL;
}

// 修复广播函数 - 避免死锁
void broadcast_message(const char *message, int exclude_sock) {
    // 复制套接字列表，避免长时间持有锁
    int target_socks[MAX_CLIENTS];
    int num_targets = 0;
    
    pthread_mutex_lock(&lock);
    for (int i = 0; i < user_count; i++) {
        if (users[i].online && users[i].sockfd != exclude_sock) {
            target_socks[num_targets++] = users[i].sockfd;
        }
    }
    pthread_mutex_unlock(&lock);
    
    // 准备消息
    ChatMessage msg;
    msg.type = BROADCAST;
    strcpy(msg.sender, "系统通知");
    strcpy(msg.message, message);
    
    // 发送消息（不持有锁）
    for (int i = 0; i < num_targets; i++) {
        send(target_socks[i], &msg, sizeof(msg), 0);
    }
}

// 保存用户数据到文件
void save_user_data() {
    FILE *file = fopen("users.dat", "wb");
    if (!file) return;
    
    fwrite(&user_count, sizeof(int), 1, file);
    for (int i = 0; i < user_count; i++) {
        fwrite(&users[i], sizeof(User), 1, file);
    }
    
    fclose(file);
}

// 从文件加载用户数据
void load_user_data() {
    FILE *file = fopen("users.dat", "rb");
    if (!file) {
        printf("未找到用户数据文件，将创建新文件\n");
        return;
    }
    
    if (fread(&user_count, sizeof(int), 1, file) != 1) {
        user_count = 0;
        fclose(file);
        return;
    }
    
    if (user_count > MAX_CLIENTS) user_count = MAX_CLIENTS;
    
    for (int i = 0; i < user_count; i++) {
        if (fread(&users[i], sizeof(User), 1, file) != 1) {
            user_count = i;
            break;
        }
    }
    fclose(file);
    
    // 初始化在线状态
    for (int i = 0; i < user_count; i++) {
        users[i].online = 0;
        users[i].sockfd = -1;
    }
    printf("成功加载 %d 个用户\n", user_count);
}